import torch
import math
import random

from torch.utils.data import Dataset



def decimal_to_list_fixed_length(n, length=5):
    # 将整数转换为字符串，并提取每一位数字
    digits = [int(digit) for digit in str(n)]
    
    # 如果位数不足，填充 0
    while len(digits) < length:
        digits.insert(0, 0)
    
    # 如果位数超过指定长度，截断列表
    return digits[:length]


# 生成训练序列
def generate_sequence_train(num_samples,seq_len,input_size,output_size):

    inputs = []
    outputs = []
    for _ in range(num_samples):
        # 随机生成一个[0, 65535]的整数
        num = random.randint(0, 65535)

        # 输入序列和输出序列
        input_sequence = []
        output_sequence = []

        # 模 2 的商和余数，直到商为 0
        while num > 0:
            remainder = num % 2
            quotient = num // 2
            input_sequence.append(decimal_to_list_fixed_length(num))  # 当前整数作为输入
            output_sequence.append((decimal_to_list_fixed_length(quotient), decimal_to_list_fixed_length(remainder)))  # 输出商和余数
            num = quotient

        # 如果输入输出序列长度不足16，则补充0
        while len(input_sequence) < 16:
            input_sequence.append(decimal_to_list_fixed_length(0))  # 输入补 0
        while len(output_sequence) < 16:
            output_sequence.append((decimal_to_list_fixed_length(0), decimal_to_list_fixed_length(0)))  # 输出补 (商为0, 余数为0)

        # 将输入输出转为 PyTorch 张量
        input_tensor = torch.tensor(input_sequence).float()
        output_tensor = torch.tensor(output_sequence).float()

        input_tensor = torch.reshape(input_tensor,[seq_len,input_size])
        output_tensor = torch.reshape(output_tensor,[seq_len,output_size])

        inputs.append(input_tensor)
        outputs.append(output_tensor)

    return inputs, outputs


# 生成测试序列
def generate_sequence_test(num_samples,seq_len,input_size,output_size):

    inputs = []
    outputs = []
    for _ in range(num_samples):
        # 随机生成一个[0, 65535]的整数
        num = random.randint(0, 65535)

        # 输入序列和输出序列
        input_sequence = []
        output_sequence = []

        input_sequence.append(decimal_to_list_fixed_length(num))  # 当前整数作为输入

        # 模 2 的商和余数，直到商为 0
        while num > 0:
            remainder = num % 2
            quotient = num // 2
            output_sequence.append( decimal_to_list_fixed_length(remainder,1))  # 输出余数
            num = quotient

        # 如果输出序列长度不足16，则补充0
        while len(output_sequence) < 16:
            output_sequence.append( decimal_to_list_fixed_length(0,1))  # 输出补 (余数为0)

        # 将输入输出转为 PyTorch 张量
        input_tensor = torch.tensor(input_sequence).float()
        output_tensor = torch.tensor(output_sequence).float()

        input_tensor = torch.reshape(input_tensor,[seq_len,input_size])
        output_tensor = torch.reshape(output_tensor,[seq_len,output_size])

        inputs.append(input_tensor)
        outputs.append(output_tensor)

    return inputs, outputs


class CustomDatasetTrain(Dataset):
    def __init__(self,num_samples,seq_len,input_size,output_size):

        self.input_tensor, self.output_tensor = generate_sequence_train(num_samples,seq_len,input_size,output_size)
    
    def __len__(self):
        return len(self.input_tensor)
    
    def __getitem__(self, idx):
        return self.input_tensor[idx], self.output_tensor[idx]


class CustomDatasetTest(Dataset):
    def __init__(self,num_samples,seq_len,input_size,output_size):

        self.input_tensor, self.output_tensor = generate_sequence_test(num_samples,seq_len,input_size,output_size)
    
    def __len__(self):
        return len(self.input_tensor)
    
    def __getitem__(self, idx):
        return self.input_tensor[idx], self.output_tensor[idx]

